import subprocess

import sys

class Logger(object):
    def __init__(self, filename='default.log', stream=sys.stdout):
        self.terminal = stream
        self.log = open(filename, 'a')

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)
        self.flush()

    def flush(self):
        self.terminal.flush()
        self.log.flush()

    def close(self):
        self.log.close()

sys.stdout = Logger('a.log', sys.stdout)
sys.stderr = Logger('a.log', sys.stderr)

# 定义要运行的Python脚本和参数
script = 'benchmark.py'
arg = '--cfg ./config/'
configs = ['face_detection_yunet.yaml',
           'face_recognition_sface.yaml',
           'facial_expression_recognition.yaml',
           'handpose_estimation_mediapipe.yaml',
           'human_segmentation_pphumanseg.yaml',
           'image_classification_mobilenet.yaml',
           'image_classification_ppresnet.yaml',
           'license_plate_detection_yunet.yaml',
           'object_detection_nanodet.yaml',
           'object_detection_yolox.yaml',
           'object_tracking_vittrack.yaml',
           'palm_detection_mediapipe.yaml',
           'person_detection_mediapipe.yaml',
           'person_reid_youtureid.yaml',
           'pose_estimation_mediapipe.yaml',
           'qrcode_wechatqrcode.yaml',
           'text_detection_ppocr.yaml',
           'text_recognition_crnn.yaml'
          ]

try:
    for config in configs:

        print('------------------------- ' + 'config: ' + config + ' -------------------------')

        # 构建命令字符串
        # ./perf_information_get.sh "python3 benchmark.py --cfg ./config/face_detection_yunet.yaml" ./result
        command = './perf_information_get.sh' + f' "python3 {script} {arg}{config}"' + ' ./result'

        try:
            # 使用 Popen 运行命令，并捕获其输出
            process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print()

            # 实时捕获并处理输出
            for line in iter(process.stdout.readline, b''):
                sys.stdout.write(line.decode())
            for line in iter(process.stderr.readline, b''):
                sys.stderr.write(line.decode())

            process.stdout.close()
            process.stderr.close()
            process.wait()

            if process.returncode != 0:
                print(f'Command failed with {process.returncode}')
                break
        except subprocess.CalledProcessError as e:
            print(f'Command failed with {e.returncode}')
            break
        except KeyboardInterrupt:
            print('Process interrupted by user')
            process.terminate()
            break
finally:
    # 关闭日志文件
    sys.stdout.close()
    sys.stderr.close()
    # 恢复原始的 sys.stdout
    sys.stdout = sys.__stdout__
    sys.stderr = sys.__stderr__